"""
llm_client.py - 通用异步 LLM 客户端封装

封装了与多种兼容 OpenAI API 格式的 LLM 服务（如阿里云、Kimi、Groq 等）交互的逻辑。
支持自动配置加载、异步上下文管理、JSON 模式输出与错误处理。
"""
import asyncio
import json
import os
from typing import Any, Dict, List, Optional, Union
from dotenv import load_dotenv
from openai import AsyncOpenAI, APIError
import yaml


class ConfigurationError(Exception):
    """当配置出现问题时抛出。"""
    pass


class LLMClient:
    """
    通用异步 LLM 客户端，支持多种兼容 OpenAI API 的提供商。

    特性：
    - 支持预定义提供商（OpenAI, Ali, Kimi, Groq, SiliconFlow）
    - 从 YAML 文件或字典加载配置
    - 自动从环境变量读取 API Key
    - 支持 JSON 模式输出并自动解析
    - 异步上下文管理器（async with）
    """

    # 预定义提供商配置
    _PREDEFINED_PROVIDERS: Dict[str, Dict[str, str]] = {
        "openai": {
            "url": "https://api.openai.com/v1",
            "api_key_env": "OPENAI_API_KEY"
        },
        "ali": {
            "url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "api_key_env": "ALI_API_KEY"
        },
        "kimi": {
            "url": "https://api.moonshot.cn/v1",
            "api_key_env": "KIMI_API_KEY"
        },
        "siliconflow": {
            "url": "https://api.siliconflow.cn/v1",
            "api_key_env": "SILICONFLOW_API_KEY"
        },
        "groq": {
            "url": "https://api.groq.com/openai/v1",
            "api_key_env": "GROQ_API_KEY"
        }
    }

    def __init__(self, config: Optional[Dict] = None, config_path: Optional[str] = None):
        """
        初始化LLM客户端。

        Args:
            config (Optional[Dict]): 直接传入配置字典。
            config_path (Optional[str]): 配置文件（.yaml）的路径。
                     e.g.
                     ```# my_config.yaml
                        llm_setting:
                        provider: "ali"  # 'ali', 'kimi', 'openai', etc.
                        model: "qwen-turbo"
                        temperature: 0.5
                        max_tokens: 4096
                        timeout: 60 # 请求超时时间（秒）
        """
        load_dotenv()  # 加载 .env 文件中的环境变量

        if config:
            self.config = config
        elif config_path:
            self.config = self._load_config_from_path(config_path)
        else:
            print("[警告] 未提供配置文件或config字典，使用默认Ali配置。")
            self.config = {
                "llm_setting": {
                    "provider": "ali",
                    "model": "qwen3-max",
                    "temperature": 0.7,
                    "top_p": 0.8,
                    "max_tokens": 1024,
                    "timeout": 30
                }
            }

        self.llm_setting = self.config.get("llm_setting")
        if not self.llm_setting:
            raise ConfigurationError("配置中缺少 'llm_setting' 部分。")

        self._client: Optional[AsyncOpenAI] = None
        self._provider_config: Optional[Dict] = None
        self._is_initialized = False

    def _load_config_from_path(self, path: str) -> Dict:
        """从YAML文件加载配置。"""
        try:
            with open(path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            raise ConfigurationError(f"配置文件未找到: {path}")
        except Exception as e:
            raise ConfigurationError(f"加载或解析配置文件失败: {e}")

    async def initialize(self):
        """
        异步初始化LLM客户端。
        根据配置验证并创建底层的AsyncOpenAI客户端。
        """
        if self._is_initialized:
            return

        provider = self.llm_setting.get("provider")
        if not provider:
            raise ConfigurationError("LLM配置中缺少 'provider'。")

        if provider not in self._PREDEFINED_PROVIDERS:
            raise ConfigurationError(f"不支持的LLM provider: '{provider}'。")

        self._provider_config = self._PREDEFINED_PROVIDERS[provider]
        api_key_env = self._provider_config["api_key_env"]
        api_key = os.getenv(api_key_env)

        if not api_key:
            raise ConfigurationError(f"缺少环境变量 '{api_key_env}'，请为 '{provider}' 设置API Key。")

        self._client = AsyncOpenAI(
            api_key=api_key,
            base_url=self._provider_config["url"],
            timeout=self.llm_setting.get("timeout", 30.0)  # 从配置中读取超时，提供默认值
        )
        self._is_initialized = True
        print(f"LLMClient 初始化成功. Provider: {provider}, Model: {self.llm_setting.get('model')}")

    async def close(self):
        """异步关闭客户端并释放资源。"""
        if self._client:
            await self._client.close()
            self._client = None
        self._is_initialized = False
        print("LLMClient 已关闭.")

    async def chat(
            self,
            system_prompt: str,
            user_prompt: str,
            as_json: bool = False,
            **kwargs
    ) -> Union[Dict[str, Any], str]:
        """
        执行一次聊天补全请求。

        Args:
            system_prompt (str): 系统提示。
            user_prompt (str): 用户提示。
            as_json (bool): 是否要求模型返回JSON格式并自动解析。
            **kwargs: 其他传递给OpenAI `chat.completions.create`的参数，
                      如 temperature, max_tokens 等。会覆盖配置文件中的默认值。

        Returns:
            Union[Dict[str, Any], str]: 如果 as_json=True，返回解析后的字典。
                                        否则，返回模型的文本响应。

        Raises:
            RuntimeError: 如果客户端未初始化。
            ValueError: 如果要求JSON但模型返回格式无效。
            APIError: 如果发生API通信错误。
        """
        if not self._is_initialized or not self._client:
            raise RuntimeError("客户端未初始化。请在使用前调用 'await client.initialize()' 或使用 'async with'。")

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

        # 合并参数：kwargs > config > defaults
        request_params = {
            "model": kwargs.get("model", self.llm_setting.get("model")),
            "temperature": kwargs.get("temperature", self.llm_setting.get("temperature", 0.7)),
            "max_tokens": kwargs.get("max_tokens", self.llm_setting.get("max_tokens", 2048)),
        }

        # 其他可能的参数（如 top_p, presence_penalty 等）也从 kwargs 覆盖
        # 过滤掉已设置的，避免重复
        for k, v in kwargs.items():
            if k not in request_params:
                request_params[k] = v

        if as_json:
            request_params["response_format"] = {"type": "json_object"}

        content = None
        try:
            response = await self._client.chat.completions.create(
                messages=messages,
                **request_params
            )

            if not response.choices:
                raise APIError("API返回了空的choices列表。", response=response, body=None)

            content = response.choices[0].message.content
            if content is None:
                raise APIError("API返回的消息内容为空。", response=response, body=None)

            if as_json:
                # 增强JSON解析的鲁棒性
                cleaned_content = content.strip()
                if cleaned_content.startswith("```json"):
                    cleaned_content = cleaned_content[7:].strip()
                if cleaned_content.endswith("```"):
                    cleaned_content = cleaned_content[:-3].strip()
                return json.loads(cleaned_content)

            return content

        except json.JSONDecodeError as e:
            error_message = f"LLM未能返回有效的JSON格式。原始内容: '{content}'"
            raise ValueError(error_message) from e
        except APIError as e:
            # 重新抛出，让上层知道是API问题
            print(f"发生API错误: {e}")
            raise
        except Exception as e:
            print(f"发生未知错误: {e}")
            raise

    # 实现异步上下文管理器协议
    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()


async def run_text_and_json_queries():
    print("--- 示例1: 使用 'async with' ---")
    try:
        # 自动处理初始化和关闭
        async with LLMClient() as client:
            # (A) 执行一个标准的文本查询
            print("\n[执行文本查询...]")
            system_prompt_text = "你是一位热情洋溢的诗人。"
            user_prompt_text = "为春天写一首五言绝句。"
            poem = await client.chat(system_prompt_text, user_prompt_text)
            print(f"诗人回答:\n{poem}")

            # (B) 执行一个JSON查询
            print("\n[执行JSON查询...]")
            system_prompt_json = "你是一个电影信息助手。根据电影名称，以JSON格式返回其导演和上映年份。JSON必须包含 'director' 和 'release_year' 两个键。"
            user_prompt_json = "《星际穿越》"
            movie_info = await client.chat(system_prompt_json, user_prompt_json, as_json=True)
            print(f"电影信息 (类型: {type(movie_info)}): {movie_info}")
            print(f"导演: {movie_info.get('director')}")

    except (ConfigurationError, ValueError, APIError) as e:
        print(f"\n发生错误: {e}")


# --- 示例2: 如何集成到您的 Agent 函数中 ---
class MyReactAgent:
    def __init__(self, llm_config_path: str):
        self.llm_config_path = llm_config_path
        # Agent的其他初始化...

    async def _internal_llm_call(self, system_prompt: str, user_prompt: str, as_json: bool = True):
        """
        一个集成了LLMClient的内部LLM调用函数。
        这就是您想要的，一个干净的集成点。
        """
        try:
            # 每次调用时创建一个新的客户端实例，通过 'async with' 管理其生命周期
            # 这种方式对于无状态的函数调用非常健壮
            async with LLMClient() as client:
                response = await client.chat(
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    as_json=as_json,
                    # 可以在这里覆盖配置中的参数
                    temperature=0.1
                )
                return response
        except Exception as e:
            print(f"Agent的LLM调用失败: {e}")
            # 根据Agent的错误处理逻辑，可以选择向上抛出或返回特定错误信息
            raise


async def run_agent_example():
    print("\n\n--- 示例2: 在Agent中集成 ---")
    agent = MyReactAgent(llm_config_path="my_config.yaml")

    # 使用Agent的内部函数
    system_prompt = ("""
        你是一位顶级的AI项目规划师和系统架构师。你的任务是分析用户的最终目标和可用的工具集，然后将其分解为一个详细、有序、可执行的子任务计划。
        
            **最终目标 (User Goal):**
           读取目录中的pdf文件，并返回摘要和大纲
        
            **可用工具 (Available Tools):**
           ls <path> 列出当前目录内容
           read_pdf <path> <max_pages>  读入1-max_pages pdf内容。
        
        
            **你的任务与规则:**
            1.  **分解任务**: 将用户的复杂目标分解成一系列更小、更简单的原子步骤。每个步骤都应该是一个清晰、独立的子任务。
            2.  **工具可行性**: 确保你规划的每一步都是可以通过一个或多个可用工具来现实地完成的。不要幻想不存在的功能。
            3.  **逻辑顺序**: 安排这些步骤的执行顺序，确保它们符合逻辑。如果步骤之间存在依赖关系，请在计划中体现出来。
            4.  **清晰描述**: 为每个步骤提供一个简洁明了的描述，说明该步骤需要完成什么。
            5.  **语言一致性**: 你的计划描述必须使用与用户目标相同的语言（例如，如果目标是中文，计划也必须是中文）。
        
            **输出格式:**
            你必须严格按照以下JSON格式输出你的计划。不要包含任何额外的解释或文本。
        
            {{
                "plan": [
                    {{
                        "step_id": 1,
                        "description": "第一步的描述...",
                        "estimated_tools": ["tool_name_1", "tool_name_2"],
                        "dependencies": []
                    }},
                    {{
                        "step_id": 2,
                        "description": "第二步的描述，它依赖于第一步的结果...",
                        "estimated_tools": ["tool_name_3"],
                        "dependencies": [1]
                    }}
                ]
            }}
        
            现在，请根据用户的最终目标和可用工具，生成你的JSON计划。
    """)
    user_prompt = ""

    city_data = await agent._internal_llm_call(system_prompt, user_prompt, as_json=True)

    if city_data:
        print(f"Agent返回: {city_data}")


if __name__ == "__main__":

    # 运行第一个示例
    asyncio.run(run_text_and_json_queries())

    # 运行第二个示例
    asyncio.run(run_agent_example())
